{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/EdjeElectronics/TensorFlow-Lite-Object-Detection-on-Android-and-Raspberry-Pi/blob/master/Train_TFLite1_Object_Detection_Model.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fF8ysCfYKgTP"
      },
      "source": [
        "# TensorFlow Lite v1 Object Detection API in Colab\n",
        "**Author:** Evan Juras, [EJ Technology Consultants](https://ejtech.io)\n",
        "\n",
        "**Last updated:** 10/12/22\n",
        "\n",
        "**GitHub:** [TensorFlow Lite Object Detection](https://github.com/EdjeElectronics/TensorFlow-Lite-Object-Detection-on-Android-and-Raspberry-Pi)\n",
        "\n",
        "\n",
        "# Introduction\n",
        "\n",
        "This notebook implements [The TensorFlow Object Detection Library](https://github.com/tensorflow/models/tree/master/research/object_detection) for training an SSD-MobileNet model using your own dataset. Note that this notebook uses TensorFlow 1 rather than TensorFlow 2, because TensorFlow 1 works better for quantizing SSD-MobileNet models. If you want to use TensorFlow 2, [clink this link to go to the TensorFlow 2 version of this notebook](https://colab.research.google.com/github/EdjeElectronics/TensorFlow-Lite-Object-Detection-on-Android-and-Raspberry-Pi/blob/master/Train_TFLite2_Object_Detction_Model.ipynb).\n",
        "\n",
        "*Note: This notebook is not maintained or updated as frequently as the TF2 notebook.*\n",
        "\n",
        "I made a YouTube video that walks through the TensorFlow 2 notebook step by step, and the process is very similar for this TensorFlow 1 notebook. Please take a look at the video if you are confused about any steps in this notebook.\n",
        "\n",
        "*Link to video to be added here*\n",
        "\n",
        "###Working with Colab\n",
        "Simply click the play button on sections of code to execute them. As the code executes, any outputs will be displayed in a block beneath the code. Once they're done executing, a green checkmark will appear next to the section to indicate it's finished running."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7EOtpvlLeS0"
      },
      "source": [
        "# 1. Install TensorFlow Object Detection Dependencies\n",
        "First, we'll install the TensorFlow Object Detection API in this Google Colab instance. This requires cloning the [TensorFlow models repository](https://github.com/tensorflow/models) and running a couple installation commands. Click the play button to run the following sections of code.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QOFtpDlci50l"
      },
      "outputs": [],
      "source": [
        "# Install CUDA 10.0 (needed for TF v1.15.3)\n",
        "!sudo apt-get update\n",
        "!sudo apt-get install cuda-10.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KeoP1s8gi8tL"
      },
      "outputs": [],
      "source": [
        "# Install cuDNN 7.6.0 (needed for TF v1.15.3)\n",
        "# Download and extract cuDNN files\n",
        "!wget https://www.dropbox.com/s/k6xqrje655q4aty/cudnn-10.0-linux-x64-v7.6.0.64.tgz\n",
        "!tar -xf cudnn-10.0-linux-x64-v7.6.0.64.tgz\n",
        "\n",
        "# Copy cuDNN libraries to appropriate folders\n",
        "!sudo cp cuda/include/cudnn*.h /usr/local/cuda-10.0/include \n",
        "!sudo cp -P cuda/lib64/libcudnn* /usr/local/cuda-10.0/lib64 \n",
        "!sudo chmod a+r /usr/local/cuda-10.0/include/cudnn*.h /usr/local/cuda-10.0/lib64/libcudnn*\n",
        "!sudo cp -P cuda/lib64/libcudnn* /usr/lib64-nvidia"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ypWGYdPlLRUN"
      },
      "outputs": [],
      "source": [
        "# Clone the tensorflow models repository from GitHub\n",
        "!git clone --depth 1 https://github.com/tensorflow/models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6QPmVBSlLTzM"
      },
      "outputs": [],
      "source": [
        "# Install the Object Detection API\n",
        "%%bash\n",
        "cd models/research/\n",
        "protoc object_detection/protos/*.proto --python_out=.\n",
        "cp object_detection/packages/tf1/setup.py .\n",
        "python -m pip install --use-feature=2020-resolver ."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kH3fwht2w17a"
      },
      "source": [
        "If you get any errors running the following code block, you can ignore them!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kDT0ctlBlXTw"
      },
      "outputs": [],
      "source": [
        "# Install TensorFlow\n",
        "!pip install tensorflow-gpu==1.15.3\n",
        "\n",
        "# A couple other packages have to be downgraded to work with the TF1 API\n",
        "!pip install numpy==1.17.4\n",
        "!pip install pycocotools==2.0.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wh_HPMOqWH9z"
      },
      "outputs": [],
      "source": [
        "# Run Model Bulider Test file, just to verify everything's working properly\n",
        "!python /content/models/research/object_detection/builders/model_builder_tf1_test.py\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IPbU4I7aL9Fl"
      },
      "source": [
        "# 2. Upload Image Dataset and Prepare Training Data\n",
        "\n",
        "In this section, we'll upload our training images and use TFRecord generation scripts to prepare the training data for TensorFlow. we'll upload our images, split them into train, validation, and test folders, and then run scripts for creating TFRecords from our data.\n",
        "\n",
        "First, on your local PC, zip all your training images and XML files into a single folder called \"images.zip\". The files should be directly inside the zip folder (wihout any additional nested folders) as shown below:\n",
        "```\n",
        "images.zip\n",
        "-- img1.jpg\n",
        "-- img1.xml\n",
        "-- img2.jpg\n",
        "-- img2.xml\n",
        "...\n",
        "```\n",
        "There are two options for moving the image files to this Colab instance: you can upload them directly, or you can copy them from your Google Drive. If you have a slow internet connection or more than 50MB worth of images, I recommend using Google Drive. Otherwise, you can just upload them through Colab. \n",
        "\n",
        "#### Option 1. Upload through Google Colab\n",
        "Upload the \"images.zip\" file to the Google Colab instance by clicking the \"Files\" icon on the left hand side of the browser, and then the \"Upload to session storage\" icon. Select the zip folder to upload it.\n",
        "\n",
        "![image.png]()\n",
        "\n",
        "Once it's uploaded, we'll run a few commands to unzip it and set up our image directories. These directories are created in the /content folder in this instance's filesystem. You can browse the filesystem by clicking the \"Files\" icon on the left.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wpi58u0-V_jN"
      },
      "source": [
        "#### Option 2. Copy from Google Drive\n",
        "You can also upload your images to your personal Google Drive, mount the drive on this Colab session, and copy them over to the Colab filesystem. This option works well if you want to upload the images beforehand so you don't have to wait for them to upload each time you restart this Colab.\n",
        "\n",
        "First, upload the \"images.zip\" file to your Google Drive, and make note of the folder you uploaded them to. Replace `MyDrive/path/to/images.zip` with the path to your zip file. (For example, I uploaded the zip file to folder called \"cat-toys1\", so I would use `MyDrive/cat-toys1/images.zip` for the path). Then, run the following block of code to mount your Google Drive to this Colab session and copy the folder to this filesystem."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tLgAPsQsfTLs"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')\n",
        "\n",
        "!cp /content/gdrive/MyDrive/path/to/images.zip /content"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DLIu4GdjxgWu"
      },
      "source": [
        "#### Option 3. Use my bird, squirrel, raccoon dataset\n",
        "I've uploaded a dataset containing 800 labeled images of birds, squirrels, and raccoons. You can use this dataset if you just want to work through the process of training a TFLite model on custom data. Run the following code block to download the dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bH6huqCBxhFb"
      },
      "outputs": [],
      "source": [
        "!wget -O /content/images.zip https://www.dropbox.com/s/en33x280e4z3wbt/BSR2.zip?dl=0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U-eXEQICx2ug"
      },
      "source": [
        "## Split images into train, validation, and test folders\n",
        "At this point, whether you used Option 1, 2, or 3, you should be able to click the folder icon on the left and see your \"images.zip\" file in the list of files. Now that the dataset is uploaded, let's unzip it and create some folders to hold the images. These directories are created in the /content folder in this instance's filesystem. You can browse the filesystem by clicking the \"Files\" icon on the left."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3n-E68MJbU2I"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/images\n",
        "!unzip -q images.zip -d /content/images/all\n",
        "!mkdir /content/images/train; mkdir /content/images/validation; mkdir /content/images/test\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n-6RIcrwbQMh"
      },
      "source": [
        "Next, we'll split the images into train, validation, and test sets. Here's what each set is used for:\n",
        "\n",
        "\n",
        "*   **Train**: These are the actual images used to train the model. In each step of training, a batch of images from the \"train\" set is passed into the neural network. The network predicts classes and locations of objects in the images. The training algorithm calculates the loss (i.e. how \"wrong\" the predictions were) and adjusts the network weights through backpropagation.\n",
        "\n",
        "\n",
        "*   **Validation**: Images from the \"validation\" set can be used by the training algorithm to check the progress of training and adjust hyperparameters (like learning rate). Unlike \"train\" images, these images are only used periodically during training (i.e. once every certain number of training steps).\n",
        "\n",
        "\n",
        "* **Test**: These images are never seen by the neural network during training. They are intended to be used by a human to perform final testing of the model to check how accurate the model is.\n",
        "\n",
        "I wrote a Python script to randomly move 80% of the images to the \"train\" folder, 10% to the \"validation\" folder, and 10% to the \"test\" folder. Click play on the following block to download the script and execute it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_8V38Uk2yBUZ"
      },
      "outputs": [],
      "source": [
        "!wget https://raw.githubusercontent.com/EdjeElectronics/TensorFlow-Lite-Object-Detection-on-Android-and-Raspberry-Pi/master/util_scripts/train_val_test_split.py\n",
        "!python train_val_test_split.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p--K1PJXEgNo"
      },
      "source": [
        "## Create TFRecords\n",
        "Finally, we need to convert the images into a data file format called TFRecords, which are used by TensorFlow for training. We'll use Python scripts to automatically convert the data into TFRecord format. Before running them, we need to define a labelmap for our classes. \n",
        "\n",
        "The code section below will create a \"labelmap.txt\" file that contains a list of classes. Replace the `class1`, `class2`, `class3` text with your own classes (for example, `bird`, `squirrel`, `raccoon`), adding a new line for each class. Then, click play to execute the code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_DE_r4MKY7ln"
      },
      "outputs": [],
      "source": [
        "### This creates a a \"labelmap.txt\" file with a list of classes the object detection model will detect.\n",
        "%%bash\n",
        "cat <<EOF >> /content/labelmap.txt\n",
        "class1\n",
        "class2\n",
        "class3\n",
        "EOF"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5pa2VYhTIT1l"
      },
      "source": [
        "Download and run the data conversion scripts by clicking play on the following two sections of code. They will create TFRecord files for the train and validation datasets, as well as a `labelmap.pbtxt` file which contains the labelmap in a different format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v3KHhgrpHved"
      },
      "outputs": [],
      "source": [
        "# Download data conversion scripts\n",
        "! wget https://raw.githubusercontent.com/EdjeElectronics/TensorFlow-Lite-Object-Detection-on-Android-and-Raspberry-Pi/master/util_scripts/create_csv.py\n",
        "! wget https://raw.githubusercontent.com/EdjeElectronics/TensorFlow-Lite-Object-Detection-on-Android-and-Raspberry-Pi/master/util_scripts/create_tfrecord.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5tdDbTmHYwu-"
      },
      "outputs": [],
      "source": [
        "# Create CSV data files and TFRecord files\n",
        "!python3 create_csv.py\n",
        "!python3 create_tfrecord.py --csv_input=images/train_labels.csv --labelmap=labelmap.txt --image_dir=images/train --output_path=train.tfrecord\n",
        "!python3 create_tfrecord.py --csv_input=images/validation_labels.csv --labelmap=labelmap.txt --image_dir=images/validation --output_path=val.tfrecord"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1XcSBULRzNZ_"
      },
      "source": [
        "We'll store the locations of the TFRecord and labelmap files as variables so we can reference them later in this Colab session."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YUd2wtfrqedy"
      },
      "outputs": [],
      "source": [
        "train_record_fname = '/content/train.tfrecord'\n",
        "val_record_fname = '/content/val.tfrecord'\n",
        "label_map_pbtxt_fname = '/content/labelmap.pbtxt'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I2MAcgJ53STW"
      },
      "source": [
        "# 3. Set Up Training Configuration \n",
        "\n",
        "In this section, we'll set up an SSD-MobileNet model training configuration. We'll specifiy which pretrained TensorFlow model we want to use from the [TensorFlow 1 Object Detection Model Zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf1_detection_zoo.md). Each model also comes with a configuration file that points to file locations, sets training parameters (such as learning rate and total number of training steps), and more. We'll modify the configuration file for our custom training job.\n",
        "\n",
        "The first section of code lists out some models availabe in the TF1 Model Zoo and defines some filenames that will be used later to download the model and config file. This makes it easy to manage which model you're using and to add other models to the list later. \n",
        "\n",
        "Set the \"chosen_model\" variable to match the name of the model you'd like to train with. It's currently set to use the \"ssd-mobilenet-v1-quantized\" model. Then, click play on the next three sections to specify and download the pretrained model file and configuration file.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gN0EUEa3e5Un"
      },
      "outputs": [],
      "source": [
        "# Change the \"chosen_model\" variable to select one of the models listed below (from the TF1 object detection zoo)\n",
        "\n",
        "chosen_model = 'ssd-mobilenet-v2-quantized'\n",
        "\n",
        "MODELS_CONFIG = {\n",
        "    'ssd-mobilenet-v1-quantized': {\n",
        "        'model_name': 'ssd_mobilenet_v1_quantized_coco',\n",
        "        'base_pipeline_file': 'ssd_mobilenet_v1_quantized_300x300_coco14_sync.config',\n",
        "        'pretrained_checkpoint': 'ssd_mobilenet_v1_quantized_300x300_coco14_sync_2018_07_18.tar.gz',\n",
        "    },\n",
        "    'ssd-mobilenet-v2-quantized': {\n",
        "        'model_name': 'ssd_mobilenet_v2_quantized_coco',\n",
        "        'base_pipeline_file': 'ssd_mobilenet_v2_quantized_300x300_coco.config',\n",
        "        'pretrained_checkpoint': 'ssd_mobilenet_v2_quantized_300x300_coco_2019_01_03.tar.gz',\n",
        "    }\n",
        "}\n",
        "\n",
        "model_name = MODELS_CONFIG[chosen_model]['model_name']\n",
        "pretrained_checkpoint = MODELS_CONFIG[chosen_model]['pretrained_checkpoint']\n",
        "base_pipeline_file = MODELS_CONFIG[chosen_model]['base_pipeline_file']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kG4TmJUVrYQ7"
      },
      "outputs": [],
      "source": [
        "# Create \"mymodel\" folder for holding pre-trained weights and configuration files\n",
        "%mkdir /content/models/mymodel/\n",
        "%cd /content/models/mymodel/\n",
        "\n",
        "# Download pretrained model file from TensorFlow Model Zoo\n",
        "import tarfile\n",
        "download_tar = 'http://download.tensorflow.org/models/object_detection/' + pretrained_checkpoint\n",
        "\n",
        "!wget {download_tar}\n",
        "tar = tarfile.open(pretrained_checkpoint)\n",
        "tar.extractall()\n",
        "tar.close()\n",
        "\n",
        "# Download base training configuration file\n",
        "download_config = 'https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/samples/configs/' + base_pipeline_file\n",
        "!wget {download_config}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3jOYxQ20wXVr"
      },
      "source": [
        "Now that we've downloaded our model and config file, we need to modify the configuration file with some high-level training parameters. The following variables are used to control training steps:\n",
        "\n",
        "* **num_steps**: The total amount of steps to use for training the model. A good number to start with is 40,000 steps. You can use more steps if you notice the loss metrics are still decreasing by the time training finishes. The more steps, the longer the training. Training can also be stopped early if loss flattens out before reaching the specified number of steps.\n",
        "* **batch_size**: The number of images to use per training step. A larger batch size allows a model to be trained in fewer steps, but the size is limited by the GPU memory available for training. For the GPUs used in Colab instances, 16 is typically a good number.\n",
        "\n",
        "* **quant_delay_steps**: (Only used for quantization-aware training) After this many steps, the training algorithm will insert \"fake\" quantization nodes in the network to simulate the effects of quantization. A good starting point is half the total number of training steps. More information is available [here](https://neuralet.com/article/quantization-of-tensorflow-object-detection-api-models/).\n",
        "\n",
        "Other training information, like the location of the pretrained model file, the config file, and total number of classes are also assigned in this step.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IKiXz33QwXyP"
      },
      "outputs": [],
      "source": [
        "# Set training parameters for the model\n",
        "num_steps = 30000\n",
        "batch_size = 16\n",
        "quant_delay_steps = 15000"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b_ki9jOqxn7V"
      },
      "outputs": [],
      "source": [
        "# Set file locations and get number of classes for config file\n",
        "pipeline_fname = '/content/models/mymodel/' + base_pipeline_file\n",
        "fine_tune_checkpoint = '/content/models/mymodel/' + model_name + '/model.ckpt'\n",
        "\n",
        "def get_num_classes(pbtxt_fname):\n",
        "    from object_detection.utils import label_map_util\n",
        "    label_map = label_map_util.load_labelmap(pbtxt_fname)\n",
        "    categories = label_map_util.convert_label_map_to_categories(\n",
        "        label_map, max_num_classes=90, use_display_name=True)\n",
        "    category_index = label_map_util.create_category_index(categories)\n",
        "    return len(category_index.keys())\n",
        "num_classes = get_num_classes(label_map_pbtxt_fname)\n",
        "print('Total classes:', num_classes)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xxpOXUTx0E-n"
      },
      "source": [
        "Finally, we'll rewrite the config file to use the training parameters we just specified. The following section of code will automatically replace the necessary parameters in the downloaded .config file and save it as our custom \"pipeline_file.config\" file."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5eA5ht3_yukT"
      },
      "outputs": [],
      "source": [
        "# Write custom configuration file by slotting our dataset, model checkpoint, and training parameters into the base pipeline file\n",
        "# TO DO: change that description\n",
        "\n",
        "import re\n",
        "\n",
        "%cd /content/models/mymodel\n",
        "print('writing custom configuration file')\n",
        "\n",
        "with open(pipeline_fname) as f:\n",
        "    s = f.read()\n",
        "with open('pipeline_file.config', 'w') as f:\n",
        "    \n",
        "    # fine_tune_checkpoint\n",
        "    s = re.sub('fine_tune_checkpoint: \".*?\"',\n",
        "               'fine_tune_checkpoint: \"{}\"'.format(fine_tune_checkpoint), s)\n",
        "    \n",
        "    # tfrecord files train and test.\n",
        "    s = re.sub(\n",
        "        '(input_path: \".*?)(PATH_TO_BE_CONFIGURED/mscoco_train.*?\")', 'input_path: \"{}\"'.format(train_record_fname), s)\n",
        "    s = re.sub(\n",
        "        '(input_path: \".*?)(PATH_TO_BE_CONFIGURED/mscoco_val)(.*?\")', 'input_path: \"{}\"'.format(val_record_fname), s)\n",
        "\n",
        "    # label_map_path\n",
        "    s = re.sub(\n",
        "        'label_map_path: \".*?\"', 'label_map_path: \"{}\"'.format(label_map_pbtxt_fname), s)\n",
        "\n",
        "    # Set training steps, num_steps\n",
        "    s = re.sub('num_steps: [0-9]+',\n",
        "               'num_steps: {}'.format(num_steps), s)\n",
        "\n",
        "    # Set training batch_size.\n",
        "    s = re.sub('batch_size: [0-9]+',\n",
        "               'batch_size: {}'.format(batch_size), s)\n",
        "    \n",
        "    # Set number of classes num_classes.\n",
        "    s = re.sub('num_classes: [0-9]+',\n",
        "               'num_classes: {}'.format(num_classes), s)\n",
        "    \n",
        "    # Fine-tune checkpoint type\n",
        "    s = re.sub(\n",
        "        'fine_tune_checkpoint_type: \"classification\"', 'fine_tune_checkpoint_type: \"{}\"'.format('detection'), s)\n",
        "    \n",
        "    # Quantization delay steps\n",
        "    s = re.sub(\n",
        "        'delay: 48000', 'delay: {}'.format(quant_delay_steps), s)\n",
        "        \n",
        "    f.write(s)\n",
        "\n",
        "%cd /content\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HEsOLOMHzBqF"
      },
      "outputs": [],
      "source": [
        "# Display the custom configuration file's contents\n",
        "!cat /content/models/mymodel/pipeline_file.config"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GMlaN3rs3zLe"
      },
      "outputs": [],
      "source": [
        "# Set the path to the custom config file, and the directory to store training checkpoints in\n",
        "pipeline_file = '/content/models/mymodel/pipeline_file.config'\n",
        "model_dir = '/content/training/'"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bLe247WqxPHw"
      },
      "source": [
        "The next block modifies the training script to save checkpoints once every 1000 training steps. Change `num_eval_steps` if you'd like to save checkpoints more or less frequently."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oi_wWkJOnc-D"
      },
      "outputs": [],
      "source": [
        "# Modify model_main.py to set evaluation and checkpoint interval\n",
        "num_eval_steps = 1000\n",
        "%cd /content/models/research/object_detection\n",
        "print('Modifying model_main.py')\n",
        "\n",
        "\n",
        "with open('model_main.py', 'r') as f:\n",
        "  data = f.read()\n",
        "\n",
        "  # Set eval steps\n",
        "  data = data.replace('config = tf_estimator.RunConfig(model_dir=FLAGS.model_dir)', \n",
        "             'config = tf.estimator.RunConfig(model_dir=FLAGS.model_dir, save_checkpoints_steps={})'.format(num_eval_steps))\n",
        "  \n",
        "  f.close()\n",
        "\n",
        "!rm /content/models/research/object_detection/model_main.py\n",
        "\n",
        "with open('model_main.py', 'w') as f:\n",
        "  f.write(data)\n",
        "\n",
        "%cd /content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7wD0tda5oqIR"
      },
      "outputs": [],
      "source": [
        "!cat /content/models/research/object_detection/model_main.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XxPj_QV43qD5"
      },
      "source": [
        "# Train Custom TF1 Object Detector\n",
        "\n",
        "We're ready to train our object detection model! Model training is performed using the \"model_main_tf2.py\" script from the TF Object Detection API. We've already defined all the parameters and arguments used by `model_main_tf2.py` in previous sections of this Colab. Training will take anywhere from 2 to 6 hours, depending on the model, batch size, and number of training steps. Just click Play on the following block to begin training!\n",
        "\n",
        "*Note: It takes a few minutes for the program to display any training messages, because it only displays logs once every 100 steps. If it seems like nothing is happening, just wait a couple minutes.*\n",
        "\n",
        "*Another note: If you want to stop training early, just click on the code block while it's running and press Ctrl+M to interrupt execution. You might have to press Ctrl+M a couple times and also click the Stop button or right-click and select \"Interrupt Execution\".*\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3mSabQWqvC2R"
      },
      "outputs": [],
      "source": [
        "!pip install tensorboard==2.8.0"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TI9iCCxoNlAL"
      },
      "outputs": [],
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir '/content/training'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tQTfZChVzzpZ"
      },
      "outputs": [],
      "source": [
        "!python /content/models/research/object_detection/model_main.py \\\n",
        "    --pipeline_config_path={pipeline_file} \\\n",
        "    --model_dir={model_dir} \\\n",
        "    --alsologtostderr \\\n",
        "    --num_train_steps={num_steps}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Vk2146Ogil3"
      },
      "source": [
        "# 5. Export Trained Inference Graph\n",
        "In the left sidebar, click the folder icon to view the files in the `/content` folder. Click the `training` folder to expand it, and find the checkpoint with the highest number (e.g. `model.ckpt-29000`). Replace \"XXXX\" in the code section below with that number."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YnSEZIzl4M10"
      },
      "outputs": [],
      "source": [
        "!mkdir /content/fine_tuned_model_lite\n",
        "output_directory = '/content/fine_tuned_model_lite'\n",
        "\n",
        "# Replace \"XXXX\" in the line below with the latest checkpoint in the /content/training directory\n",
        "last_model_path = '/content/training/model.ckpt-30000'\n",
        "print(last_model_path)\n",
        "\n",
        "!python /content/models/research/object_detection/export_tflite_ssd_graph.py \\\n",
        "    --trained_checkpoint_prefix {last_model_path} \\\n",
        "    --output_directory {output_directory} \\\n",
        "    --pipeline_config_path {pipeline_file}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IwoGtqdejiez"
      },
      "source": [
        "# 6. Convert Model to TensorFlow Lite Format"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PW7TpzLmhTbW"
      },
      "outputs": [],
      "source": [
        "localpb = '/content/fine_tuned_model_lite/tflite_graph.pb'\n",
        "tflite_file = '/content/fine_tuned_model_lite/detect.tflite'\n",
        "\n",
        "# Can use Netron app to determine the input size: https://www.electronjs.org/apps/netron \n",
        "input_size = [1, 300, 300, 3]\n",
        "input_shape = {\"normalized_input_image_tensor\" : input_size}\n",
        "\n",
        "print(input_shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ec8LzJLbr-SR"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "converter = tf.lite.TFLiteConverter.from_frozen_graph(\n",
        "    localpb, \n",
        "    [\"normalized_input_image_tensor\"],\n",
        "    [\"TFLite_Detection_PostProcess\",\"TFLite_Detection_PostProcess:1\",\"TFLite_Detection_PostProcess:2\",\"TFLite_Detection_PostProcess:3\"],\n",
        "    input_shape\n",
        ")\n",
        "\n",
        "converter.allow_custom_ops = True\n",
        "converter.inference_type = tf.uint8  # Tell converter to create a quantized model\n",
        "converter.quantized_input_stats = {'normalized_input_image_tensor': (128, 128)}\n",
        "\n",
        "tflite_model = converter.convert()\n",
        "\n",
        "open(tflite_file,'wb').write(tflite_model)\n",
        "\n",
        "# Test that the conversion was successful - if any errors are generated, something went wrong!\n",
        "interpreter = tf.lite.Interpreter(model_content=tflite_model)\n",
        "interpreter.allocate_tensors()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "awZMQGVqMpVL"
      },
      "outputs": [],
      "source": [
        "!cp /content/labelmap.txt /content/fine_tuned_model_lite\n",
        "!cp /content/labelmap.pbtxt /content/fine_tuned_model_lite\n",
        "!zip -r /content/fine_tuned_model_lite.zip /content/fine_tuned_model_lite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w41-KMn7KV-T"
      },
      "source": [
        "Now, you can download the \"detect.tflite\" and \"labelmap.txt\" files and move them to your Raspberry Pi."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bQyjuMAt1WEL"
      },
      "source": [
        "# (Optional) Test out TensorFlow Lite Model\n",
        "Now that we've got our custom detection model trained and converted to TFLite format, let's try it out on some images. We'll use the images from our \"test\" folder. These images weren't seen at all by the model during training, so they'll give a faithful depiction of how accurate the model actually is.\n",
        "\n",
        "The following code is used to import the TensorFlow Lite runtime, load it into memory, run inferencing on the test images, and display the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v1IBMfrD2GEQ"
      },
      "outputs": [],
      "source": [
        "# Copied from my TFLite repository, and stripped unneeded stuff out\n",
        "\n",
        "# Import packages\n",
        "import os\n",
        "import cv2\n",
        "import numpy as np\n",
        "import sys\n",
        "import glob\n",
        "import importlib.util\n",
        "from tensorflow.lite.python.interpreter import Interpreter\n",
        "\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "%matplotlib inline\n",
        "\n",
        "PATH_TO_IMAGES='/content/images/test'\n",
        "PATH_TO_MODEL='/content/fine_tuned_model_lite/detect.tflite'\n",
        "PATH_TO_LABELS='/content/labelmap.txt'\n",
        "IMAGES_TO_TEST = 10\n",
        "min_conf_threshold=0.5\n",
        "\n",
        "# Grab all image filenames\n",
        "images = glob.glob(PATH_TO_IMAGES + '/*.jpg')\n",
        "images = images + glob.glob(PATH_TO_IMAGES + '/*.JPG')\n",
        "images = images + glob.glob(PATH_TO_IMAGES + '/*.png')\n",
        "\n",
        "\n",
        "# Load the label map into memory\n",
        "with open(PATH_TO_LABELS, 'r') as f:\n",
        "    labels = [line.strip() for line in f.readlines()]\n",
        "\n",
        "\n",
        "# Load the Tensorflow Lite model into memory\n",
        "interpreter = Interpreter(model_path=PATH_TO_MODEL)\n",
        "interpreter.allocate_tensors()\n",
        "\n",
        "# Get model details\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "height = input_details[0]['shape'][1]\n",
        "width = input_details[0]['shape'][2]\n",
        "\n",
        "floating_model = (input_details[0]['dtype'] == np.float32)\n",
        "\n",
        "input_mean = 127.5\n",
        "input_std = 127.5\n",
        "\n",
        "# Loop over every image and perform detection\n",
        "for image_path in images[0:IMAGES_TO_TEST]:\n",
        "\n",
        "    # Load image and resize to expected shape [1xHxWx3]\n",
        "    image = cv2.imread(image_path)\n",
        "    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "    imH, imW, _ = image.shape \n",
        "    image_resized = cv2.resize(image_rgb, (width, height))\n",
        "    input_data = np.expand_dims(image_resized, axis=0)\n",
        "\n",
        "    # Normalize pixel values if using a floating model (i.e. if model is non-quantized)\n",
        "    if floating_model:\n",
        "        input_data = (np.float32(input_data) - input_mean) / input_std\n",
        "\n",
        "    # Perform the actual detection by running the model with the image as input\n",
        "    interpreter.set_tensor(input_details[0]['index'],input_data)\n",
        "    interpreter.invoke()\n",
        "\n",
        "    # Retrieve detection results\n",
        "    boxes = interpreter.get_tensor(output_details[0]['index'])[0] # Bounding box coordinates of detected objects\n",
        "    classes = interpreter.get_tensor(output_details[1]['index'])[0] # Class index of detected objects\n",
        "    scores = interpreter.get_tensor(output_details[2]['index'])[0] # Confidence of detected objects\n",
        "    #num = interpreter.get_tensor(output_details[3]['index'])  # Total number of detected objects (inaccurate and not needed)\n",
        "\n",
        "    # Loop over all detections and draw detection box if confidence is above minimum threshold\n",
        "    for i in range(len(scores)):\n",
        "        if ((scores[i] > min_conf_threshold) and (scores[i] <= 1.0)):\n",
        "\n",
        "            # Get bounding box coordinates and draw box\n",
        "            # Interpreter can return coordinates that are outside of image dimensions, need to force them to be within image using max() and min()\n",
        "            ymin = int(max(1,(boxes[i][0] * imH)))\n",
        "            xmin = int(max(1,(boxes[i][1] * imW)))\n",
        "            ymax = int(min(imH,(boxes[i][2] * imH)))\n",
        "            xmax = int(min(imW,(boxes[i][3] * imW)))\n",
        "            \n",
        "            cv2.rectangle(image, (xmin,ymin), (xmax,ymax), (10, 255, 0), 2)\n",
        "\n",
        "            # Draw label\n",
        "            object_name = labels[int(classes[i])] # Look up object name from \"labels\" array using class index\n",
        "            label = '%s: %d%%' % (object_name, int(scores[i]*100)) # Example: 'person: 72%'\n",
        "            labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2) # Get font size\n",
        "            label_ymin = max(ymin, labelSize[1] + 10) # Make sure not to draw label too close to top of window\n",
        "            cv2.rectangle(image, (xmin, label_ymin-labelSize[1]-10), (xmin+labelSize[0], label_ymin+baseLine-10), (255, 255, 255), cv2.FILLED) # Draw white box to put label text in\n",
        "            cv2.putText(image, label, (xmin, label_ymin-7), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 0), 2) # Draw label text\n",
        "\n",
        "    # All the results have been drawn on the image, now display the image\n",
        "    image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)\n",
        "    plt.figure(figsize=(12,16))\n",
        "    plt.imshow(image)\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XFsuasvxFHo8"
      },
      "source": [
        "## (Optional) Compile Model for Edge TPU\n",
        "\n",
        "Now that the model has been converted to TFLite and quantized, we can compile it to run on Edge TPU devices like the [Coral USB Accelerator](https://coral.ai/products/accelerator/) or the [Coral Dev Board](https://coral.ai/products/dev-board/). This allows the model to run much faster! For more information on the Edge TPU, see my [TensorFlow Lite repository on GitHub](fine_tuned_model_lite).\n",
        "\n",
        "First, install the Edge TPU Compiler package inside this Colab instance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mUd_SNC0JSq0"
      },
      "outputs": [],
      "source": [
        "! curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -\n",
        "! echo \"deb https://packages.cloud.google.com/apt coral-edgetpu-stable main\" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list\n",
        "! sudo apt-get update\n",
        "! sudo apt-get install edgetpu-compiler\t"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "usfmdtSiJuuC"
      },
      "source": [
        "Next, compile the model. (If your model has a different filename than \"detect.tflite\", use that instead.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WTboEAWuJ0ku"
      },
      "outputs": [],
      "source": [
        "%cd /content/fine_tuned_model_lite\n",
        "!edgetpu_compiler detect.tflite\n",
        "!mv detect_edgetpu.tflite edgetpu.tflite"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oqGy2FgzKomN"
      },
      "source": [
        "The compiled model will be output in the /content folder as \"detect_edgetpu.tflite\". It gets renamed to \"edgetpu.tflite\" to be consistent with my code. Download this file using the command below. (If your model has a different filename than \"edgetpu.tflite\", use that instead.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GgvZC_y5ESYq"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!cp labelmap.txt fine_tuned_model_lite\n",
        "!cp labelmap.pbtxt fine_tuned_model_lite\n",
        "!zip -r fine_tuned_model_lite.zip fine_tuned_model_lite"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AmjqvKuuK8ZR"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download('edgetpu.tflite')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ptwpBBEWLfuJ"
      },
      "source": [
        "Now you're all set to use the Coral model! For instructions on how to run an object detection model on the Raspberry Pi using the Coral USB Acclerator, please see my video, [\"How to Use the Coral USB Accelerator with the Raspberry Pi\"](https://www.youtube.com/watch?v=qJMwNHQNOVU)."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}